{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ADAPT-QAOA algorithm\n",
    "\n",
    "In this tutorial, we explain the ADAPT-QAOA algorithm introduced in this [paper](https://journals.aps.org/prresearch/pdf/10.1103/PhysRevResearch.4.033029) and show how to implement the algorithm in cudaq.\n",
    "\n",
    "In QAOA (explained in this [Tutorial](https://nvidia.github.io/cuda-quantum/latest/applications/python/qaoa.html)), the variational Ansatz consists of $p$ layers, each containing the cost Hamiltonian $H_C$ and a mixer, $H_M$. \n",
    "$$ \\ket{\\psi_p (\\beta, \\gamma)}  = \\left ( \\prod_{k=1} ^p [e^{-i H_M \\beta_k}  e^{-i H_C \\gamma_k}]  \\right ) \\ket{\\psi_{\\mathrm{ref}}} $$\n",
    "\n",
    "where $\\ket{\\psi_{\\mathrm{ref}}} =  \\ket{+}^{\\otimes n}$, $n$ is the number of qubits and $\\gamma, \\beta$ are set of variational parameters. If these parameters are chosen such that $$\\braket{\\psi_p (\\beta, \\gamma)| H_C |\\psi_p (\\beta, \\gamma) }$$ is minimized, then the resulting energy and state provide an approximate solution to the optimization problem encoded in $H_C$. In the standard QAOA Ansatz, the mixer is chosen to be a single-qubit X rotation applied to all qubits.\n",
    "\n",
    "In ADAPT-QAOA, the single, fixed mixer $H_M$ is replaced by a set of mixers $A_k$ that change from one layer to the next.\n",
    "$$ \\ket{\\psi_p (\\beta, \\gamma) }  = \\left ( \\prod_{k=1} ^p [e^{-i A_k \\beta_k}  e^{-i H_C \\gamma_k}]  \\right ) \\ket{\\psi_{\\mathrm{ref}}} $$\n",
    "The ADAPT-QAOA algorithm can be summarized in three steps:\n",
    "\n",
    "1- Define the operator set $A_j$ (called the mixer pool, and where $A_j = A^{\\dagger}_j$ and select a suitable reference state to be the initial\n",
    "state: $\\ket{ \\psi (0) }= \\ket{\\psi_{\\mathrm{ref}}}$; $\\ket{\\psi_{\\mathrm{ref}}} =  \\ket{+}^{\\otimes n}$  is chosen as in the standard QAOA\n",
    "\n",
    "2- Prepare the current Ansatz $\\ket{\\psi(k-1)}$ and measure the energy gradient with respect to the pool, the $j$th component of which is given by $−i \\braket{\\psi(k-1)|e^{i H_C \\gamma_k} [H_C,A_j] e^{−i H_C \\gamma_k} |\\psi(k-1)}$, where the new variational parameter $\\gamma_k$ is set to a predefined value $\\gamma_0$. For the measurement, we can decompose the commutator into linear combinations of Pauli strings and measure the expectation values of the strings using general variational quantum algorithm methods. If the norm of the gradient is below a predefined threshold, then the algorithm stops, and the current state and energy estimate approximate the desired solution. If the gradient threshold is not met, modify the Ansatz (k) by adding the operator, A(k) , associated with the largest max component of the gradient: $\\ket{\\psi(k)}= e^ {−i A_\\mathrm{max} \\beta_k} e^{−i H_C \\gamma_k} \\ket{\\psi(k-1)}$, where $\\beta_k$ is a new variational parameter\n",
    "\n",
    "3- Optimize all parameters currently in the Ansatz $\\beta_m, \\gamma_m = 1, 2, ...k$ such that $\\braket{\\psi (k)|H_C|\\psi(k)}$ is minimized, and return to the second step.\n",
    "\n",
    "\n",
    "Below is a schematic representation of the ADAPT-QAOA algorithm explained above.\n",
    "\n",
    "<div>\n",
    "<img src=\"images/adapt-qaoa.png\" width=\"1000\">\n",
    "</div>\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cudaq\n",
    "from cudaq import spin\n",
    "from scipy.optimize import minimize\n",
    "import random\n",
    "\n",
    "cudaq.set_target(\"nvidia\", option=\"fp64\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simulation input:\n",
    "\n",
    "The gradients used to select new operators are sensitive to the initial values for $\\gamma$. We initialize the parameter at $\\gamma_0 =0.01$ as explained in the paper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The Max-cut graph:\n",
    "\n",
    "# Example 1\n",
    "#true_energy=-12.0\n",
    "#nodes = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]\n",
    "#edges = [[0,4], [1,4], [0,6], [3,6], [5,6], [2,7], [1,8], [3,8], [5,8], [6,8], [7,8], [2,9], [8,9]]\n",
    "#weight = [1.0] * len(edges)\n",
    "\n",
    "# Examle 2:\n",
    "#true_energy=-5.0\n",
    "nodes = [0, 1, 2, 3, 4]\n",
    "edges = [[0, 1], [1, 2], [2, 3], [3, 0], [2, 4], [3, 4]]\n",
    "weight = [1.0] * len(edges)\n",
    "\n",
    "# Define the number of qubits.\n",
    "\n",
    "qubits_num = len(nodes)\n",
    "\n",
    "# Predefined values to terminate iteration. These can be changed based on user preference\n",
    "E_prev = 0.0       \n",
    "threshold = 1e-3   # This is a predefined value to stop the calculation. if norm of the gradient smaller than this value, calculation will be terminated.\n",
    "e_stop = 1e-7      # This is a predefined value to stop the calculation. if dE <= e_stop, calculation stop. dE=current_erg-prev_erg\n",
    "\n",
    "# Initiate the parameters gamma of the problem Hamiltonian for gradient calculation \n",
    "\n",
    "init_gamma = 0.01\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The problem Hamiltonian $H_C$ of the max-cut graph:\n",
    "\n",
    "$$ H_C= - \\frac{1}{2} \\sum_{i,j} \\omega_{i,j} (I - Z_i Z_j) $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def spin_ham (edges, weight):\n",
    "    \n",
    "    ham = 0\n",
    "\n",
    "    count = 0\n",
    "    for edge in range(len(edges)):\n",
    "\n",
    "        qubitu = edges[edge][0]\n",
    "        qubitv = edges[edge][1]\n",
    "        # Add a term to the Hamiltonian for the edge (u,v)\n",
    "        ham += 0.5 * (weight[count] * cudaq.spin.z(qubitu) * cudaq.spin.z(qubitv) -\n",
    "                            weight[count] * cudaq.spin.i(qubitu) * cudaq.spin.i(qubitv))\n",
    "        count += 1\n",
    "\n",
    "    return ham\n",
    "\n",
    "# Collect coefficients from a spin operator so we can pass them to a kernel\n",
    "def term_coefficients(ham: cudaq.SpinOperator) -> list[complex]:\n",
    "    result = []\n",
    "    for term in ham:\n",
    "        result.append(term.evaluate_coefficient())\n",
    "    return result\n",
    "\n",
    "# Collect Pauli words from a spin operator so we can pass them to a kernel\n",
    "def term_words(ham: cudaq.SpinOperator) -> list[str]:\n",
    "    # Our kernel uses these words to apply exp_pauli to the entire state.\n",
    "    # we hence ensure that each pauli word covers the entire space.\n",
    "    \n",
    "    result = []\n",
    "    for term in ham:\n",
    "        result.append(term.get_pauli_word(qubits_num))\n",
    "    return result\n",
    "\n",
    "# Generate the Spin Hamiltonian:\n",
    "\n",
    "ham = spin_ham(edges, weight)\n",
    "ham_coef = term_coefficients(ham)\n",
    "ham_word = term_words(ham)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Th operator pool $A_j$:\n",
    "\n",
    "Here, the mixer pool consists of qaoa mixer, single-qubit, and multi-qubit entangling gates. Note that more single and multi-qubit gates can be added based on the problem. Because of symmetry, we only includes the mixer pools that will have nonzero gradients as explained in the [paper](https://journals.aps.org/prresearch/pdf/10.1103/PhysRevResearch.4.033029)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def qaoa_mixer(n):\n",
    "\n",
    "    term = spin.x(0)\n",
    "\n",
    "    for i in range(1, n):\n",
    "        term = term + spin.x(i)\n",
    "\n",
    "    pool = [term]\n",
    "    return pool\n",
    "\n",
    "def qaoa_single_x(n):\n",
    "\n",
    "    pool = []\n",
    "\n",
    "    for i in range(n):\n",
    "        pool.append(cudaq.SpinOperator(spin.x(i)))\n",
    "\n",
    "    return pool\n",
    "\n",
    "def qaoa_double_ops(n):\n",
    "\n",
    "    pool = []\n",
    "\n",
    "    for i in range(n-1):\n",
    "        for j in range(i+1, n):\n",
    "            pool.append(cudaq.SpinOperator(spin.x(i)) * cudaq.SpinOperator(spin.x(j)))\n",
    "            pool.append(cudaq.SpinOperator(spin.y(i)) * cudaq.SpinOperator(spin.y(j)))\n",
    "            pool.append(cudaq.SpinOperator(spin.y(i)) * cudaq.SpinOperator(spin.z(j)))\n",
    "            pool.append(cudaq.SpinOperator(spin.z(i)) * cudaq.SpinOperator(spin.y(j)))\n",
    "\n",
    "    return pool\n",
    "\n",
    "def mixer_pool(qubits_num):\n",
    "    return (\n",
    "          qaoa_single_x(qubits_num)\n",
    "        + qaoa_mixer(qubits_num)\n",
    "        + qaoa_double_ops(qubits_num)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of pool operator:  46\n"
     ]
    }
   ],
   "source": [
    "# Generate the mixer pool\n",
    "\n",
    "pools = mixer_pool(qubits_num)\n",
    "print('Number of pool operator: ', len(pools))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The commutator $[H_C,A_j]$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate the commutator [H,Ai]\n",
    "\n",
    "def commutator(pools, ham):\n",
    "    com_op = []\n",
    "\n",
    "    for i in range(len(pools)):\n",
    "    #for i in range(2):\n",
    "        op = pools[i]\n",
    "        com_op.append(ham * op - op * ham)\n",
    "\n",
    "    return com_op\n",
    "\n",
    "grad_op = commutator(pools, ham)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "###########################################\n",
    "# Get the initial state (psi_ref)\n",
    "\n",
    "@cudaq.kernel\n",
    "def initial_state(qubits_num:int):\n",
    "    qubits = cudaq.qvector(qubits_num)\n",
    "    h(qubits)\n",
    "\n",
    "state = cudaq.get_state(initial_state, qubits_num)\n",
    "\n",
    "#print(state)\n",
    "###############################################\n",
    "\n",
    "# Circuit to compute the energy gradient with respect to the pool\n",
    "@cudaq.kernel\n",
    "def grad(state:cudaq.State, ham_word:list[cudaq.pauli_word], ham_coef: list[complex], init_gamma:float):\n",
    "    q = cudaq.qvector(state)\n",
    "\n",
    "    for i in range(len(ham_coef)):\n",
    "        exp_pauli(init_gamma * ham_coef[i].real, q, ham_word[i])\n",
    "\n",
    "\n",
    "# The qaoa circuit using the selected pool operator with max gradient\n",
    "\n",
    "@cudaq.kernel\n",
    "def kernel_qaoa(qubits_num:int, ham_word:list[cudaq.pauli_word], ham_coef:list[complex],\n",
    "        mixer_pool:list[list[cudaq.pauli_word]], gamma:list[float], beta:list[float], num_layer:int):\n",
    "\n",
    "    qubits = cudaq.qvector(qubits_num)\n",
    "\n",
    "    h(qubits)\n",
    "\n",
    "    \n",
    "    for p in range(num_layer):\n",
    "        for i in range(len(ham_coef)):\n",
    "            exp_pauli(gamma[p] * ham_coef[i].real, qubits, ham_word[i])\n",
    "        \n",
    "        for word in mixer_pool[p]:\n",
    "            exp_pauli(beta[p], qubits, word)\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Beginning of ADAPT-QAOA iteration:\n",
    "\n",
    "Note that here in this tutorial, we set the seed for the random number generator. This ensures that the random choices are reproducible in each step of iteration when choosing operator pools if more than one operator pool has the same value. For production, users need to remove that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step:  1\n",
      "Norm of the gradient:  3.468398683655509\n",
      "Mixer pool at step 1\n",
      "[['IIZIY']]\n",
      "Number of layers:  1\n",
      "Result from the step  1\n",
      "Optmized Energy:  -3.4999999999998064\n",
      "dE= : 3.4999999999998064\n",
      "\n",
      "\n",
      "Step:  2\n",
      "Norm of the gradient:  2.4501630553527365\n",
      "Mixer pool at step 2\n",
      "[['IIZIY'], ['ZIIYI']]\n",
      "Number of layers:  2\n",
      "Result from the step  2\n",
      "Optmized Energy:  -3.999999999998128\n",
      "dE= : 0.4999999999983218\n",
      "\n",
      "\n",
      "Step:  3\n",
      "Norm of the gradient:  1.99990001175888\n",
      "Mixer pool at step 3\n",
      "[['IIZIY'], ['ZIIYI'], ['ZYIII']]\n",
      "Number of layers:  3\n",
      "Result from the step  3\n",
      "Optmized Energy:  -4.499999999929474\n",
      "dE= : 0.49999999993134603\n",
      "\n",
      "\n",
      "Step:  4\n",
      "Norm of the gradient:  0.014141961855183087\n",
      "Mixer pool at step 4\n",
      "[['IIZIY'], ['ZIIYI'], ['ZYIII'], ['IIXIX']]\n",
      "Number of layers:  4\n",
      "Result from the step  4\n",
      "Optmized Energy:  -4.999999999996581\n",
      "dE= : 0.5000000000671072\n",
      "\n",
      "\n",
      "Step:  5\n",
      "Norm of the gradient:  5.778820729079865e-06\n",
      "\n",
      " Final Result:  \n",
      "\n",
      "Final mixer_pool:  [['IIZIY'], ['ZIIYI'], ['ZYIII'], ['IIXIX']]\n",
      "Number of layers:  4\n",
      "Number of mixer pool in each layer:  [1, 1, 1, 1]\n",
      "Final Energy:  -4.999999999996581\n"
     ]
    }
   ],
   "source": [
    "# Begining od ADAPT-QAOA iteration\n",
    "\n",
    "beta = []\n",
    "gamma = []\n",
    "\n",
    "mixer_pool = []\n",
    "mixer_pool_str = []\n",
    "layer=[]\n",
    "\n",
    "istep = 1\n",
    "while True:\n",
    "\n",
    "    print('Step: ', istep)\n",
    "\n",
    "    #####################\n",
    "    # compute the gradient and find the mixer pool with large values. \n",
    "    # If norm is below the predefined threshold, stop calculation\n",
    "\n",
    "    gradient_vec = []\n",
    "    for op in grad_op:\n",
    "        op = op * -1j\n",
    "        gradient_vec.append(cudaq.observe(grad, op , state, ham_word, ham_coef, init_gamma).expectation())\n",
    "        \n",
    "    # Compute the norm of the gradient vector\n",
    "    norm = np.linalg.norm(np.array(gradient_vec))\n",
    "    print('Norm of the gradient: ', norm)\n",
    "    \n",
    "    if norm <= threshold:\n",
    "        print('\\n', 'Final Result: ', '\\n')\n",
    "        print('Final mixer_pool: ', mixer_pool_str)\n",
    "        print('Number of layers: ', len(layer))\n",
    "        print('Number of mixer pool in each layer: ', layer)\n",
    "        print('Final Energy: ', result_vqe.fun)\n",
    "        \n",
    "        break\n",
    "    \n",
    "    else:\n",
    "        temp_pool = []\n",
    "        tot_pool = 0\n",
    "        \n",
    "        max_grad = np.max(np.abs(gradient_vec))\n",
    "        \n",
    "        for i in range(len(pools)):\n",
    "            if np.abs(gradient_vec[i]) >= max_grad:\n",
    "                tot_pool += 1\n",
    "                temp_pool.append(pools[i])\n",
    "        \n",
    "        # Set the seed for the random number generator\n",
    "        # This ensures that the random choices are reproducible\n",
    "        # in each step of the iteration.\n",
    "        random.seed(42)   \n",
    "        \n",
    "        layer.append(1) \n",
    "        random_mixer = random.choice(temp_pool)\n",
    "        \n",
    "        pool_added = []\n",
    "        pool_added_str = []\n",
    "        for term in random_mixer:\n",
    "            pool_added.append(cudaq.pauli_word(term.get_pauli_word(qubits_num)))\n",
    "            pool_added_str.append(term.get_pauli_word(qubits_num))\n",
    "            \n",
    "        mixer_pool.append(pool_added)\n",
    "        mixer_pool_str.append(pool_added_str)\n",
    "        \n",
    "    \n",
    "        print('Mixer pool at step', istep)\n",
    "        print(mixer_pool_str)\n",
    "        \n",
    "        num_layer = len(layer)\n",
    "        print('Number of layers: ', num_layer)\n",
    "        \n",
    "        beta_count = layer[num_layer-1]\n",
    "        init_beta = [0.0] * beta_count\n",
    "        beta = beta + init_beta\n",
    "        gamma = gamma + [init_gamma]\n",
    "        theta = gamma + beta\n",
    "        \n",
    "        def cost(theta):\n",
    "\n",
    "            theta = theta.tolist()\n",
    "            gamma = theta[:num_layer]\n",
    "            beta = theta[num_layer:]\n",
    "\n",
    "            \n",
    "            energy = cudaq.observe(kernel_qaoa, ham, qubits_num, ham_word, ham_coef,\n",
    "                                    mixer_pool, gamma, beta, num_layer).expectation()\n",
    "            #print(energy)\n",
    "            return energy\n",
    "        \n",
    "        result_vqe = minimize(cost, theta, method='BFGS', jac='2-point', tol=1e-5)\n",
    "        print('Result from the step ', istep)\n",
    "        print('Optmized Energy: ', result_vqe.fun)\n",
    "        \n",
    "        dE = np.abs(result_vqe.fun - E_prev)\n",
    "        E_prev = result_vqe.fun\n",
    "        theta = result_vqe.x.tolist()\n",
    "        erg = result_vqe.fun\n",
    "        \n",
    "        print('dE= :', dE)\n",
    "        gamma=theta[:num_layer]\n",
    "        beta=theta[num_layer:]\n",
    "        \n",
    "        if dE <= e_stop:\n",
    "            print('\\n', 'Final Result: ', '\\n')\n",
    "            print('Final mixer_pool: ', mixer_pool_str)\n",
    "            print('Number of layers: ', len(layer))\n",
    "            print('Number of mixer pool in each layer: ', layer)\n",
    "            print('Final Energy= ', erg)\n",
    "\n",
    "            break\n",
    "        \n",
    "        else:\n",
    "\n",
    "            # Compute the state of this current step for the gradient\n",
    "            state = cudaq.get_state(kernel_qaoa, qubits_num, ham_word, ham_coef,mixer_pool, gamma, beta, num_layer)\n",
    "            #print('State at step ', istep)\n",
    "            #print(state)\n",
    "            istep+=1\n",
    "            print('\\n')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " Sampling the Final ADAPT QAOA circuit \n",
      "\n",
      "The most probable max cut:  01011\n",
      "All bitstring from circuit sampling:  { 01011:2519 10100:2481 }\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print('\\n', 'Sampling the Final ADAPT QAOA circuit', '\\n')\n",
    "# Sample the circuit\n",
    "count=cudaq.sample(kernel_qaoa, qubits_num, ham_word, ham_coef,mixer_pool, gamma, beta, num_layer, shots_count=5000)\n",
    "print('The most probable max cut: ', count.most_probable())\n",
    "print('All bitstring from circuit sampling: ', count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CUDA-Q Version cu12-latest (https://github.com/NVIDIA/cuda-quantum 615d91910310605c83ea59f6afe6e7ae6dfecd28)\n"
     ]
    }
   ],
   "source": [
    "print(cudaq.__version__)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
